"""
Python AST parser for code analysis.

Provides comprehensive Python AST parsing and extraction of code elements
including functions, classes, imports, and structural information.
"""

import ast
from typing import Dict, List, Any, Optional, Set, Union
from pathlib import Path

from ...core.error_handler import AnalysisError, log_info


class PythonParser:
    """Python AST parser for code analysis."""
    
    def __init__(self):
        """Initialize the Python parser."""
        self.current_file = None
        self.current_scope_stack = []
    
    def parse_file(self, ast_tree: ast.AST, file_path: str) -> Dict[str, Any]:
        """Parse a Python file and extract all relevant information."""
        self.current_file = file_path
        self.current_scope_stack = []
        
        visitor = PythonASTVisitor()
        visitor.visit(ast_tree)
        
        return {
            'file_path': file_path,
            'functions': visitor.functions,
            'classes': visitor.classes,
            'imports': visitor.imports,
            'global_variables': visitor.global_variables,
            'function_calls': visitor.function_calls,
            'method_calls': visitor.method_calls,
            'attribute_access': visitor.attribute_access,
            'constants': visitor.constants,
            'decorators': visitor.decorators,
            'docstrings': visitor.docstrings,
            'scope_tree': visitor.scope_tree
        }


class PythonASTVisitor(ast.NodeVisitor):
    """AST visitor for extracting Python code elements."""
    
    def __init__(self):
        """Initialize the AST visitor."""
        self.functions = []
        self.classes = []
        self.imports = []
        self.global_variables = []
        self.function_calls = []
        self.method_calls = []
        self.attribute_access = []
        self.constants = []
        self.decorators = []
        self.docstrings = []
        self.scope_tree = []
        
        # Internal state
        self.current_scope = None
        self.scope_stack = []
        self.current_class = None
        self.current_function = None
    
    def _enter_scope(self, scope_type: str, name: str, node: ast.AST):
        """Enter a new scope."""
        scope = {
            'type': scope_type,
            'name': name,
            'line_start': node.lineno,
            'line_end': getattr(node, 'end_lineno', node.lineno),
            'col_start': node.col_offset,
            'col_end': getattr(node, 'end_col_offset', node.col_offset),
            'parent': self.current_scope,
            'children': []
        }
        
        if self.current_scope:
            self.current_scope['children'].append(scope)
        else:
            self.scope_tree.append(scope)
        
        self.scope_stack.append(self.current_scope)
        self.current_scope = scope
        return scope
    
    def _exit_scope(self):
        """Exit the current scope."""
        if self.scope_stack:
            self.current_scope = self.scope_stack.pop()
    
    def _get_docstring(self, node: Union[ast.FunctionDef, ast.ClassDef, ast.Module]) -> Optional[str]:
        """Extract docstring from a node."""
        if (node.body and 
            isinstance(node.body[0], ast.Expr) and 
            isinstance(node.body[0].value, ast.Constant) and 
            isinstance(node.body[0].value.value, str)):
            return node.body[0].value.value
        return None
    
    def _extract_args(self, args: ast.arguments) -> List[Dict[str, Any]]:
        """Extract function arguments information."""
        arguments = []
        
        # Regular arguments
        for i, arg in enumerate(args.args):
            arg_info = {
                'name': arg.arg,
                'type': 'positional',
                'annotation': self._get_annotation(arg.annotation) if arg.annotation else None,
                'default': None,
                'line': arg.lineno,
                'col': arg.col_offset
            }
            
            # Check for default values
            defaults_start = len(args.args) - len(args.defaults)
            if i >= defaults_start:
                default_idx = i - defaults_start
                arg_info['default'] = self._ast_to_string(args.defaults[default_idx])
            
            arguments.append(arg_info)
        
        # *args
        if args.vararg:
            arguments.append({
                'name': args.vararg.arg,
                'type': 'vararg',
                'annotation': self._get_annotation(args.vararg.annotation) if args.vararg.annotation else None,
                'line': args.vararg.lineno,
                'col': args.vararg.col_offset
            })
        
        # Keyword-only arguments
        for i, arg in enumerate(args.kwonlyargs):
            arg_info = {
                'name': arg.arg,
                'type': 'keyword_only',
                'annotation': self._get_annotation(arg.annotation) if arg.annotation else None,
                'default': None,
                'line': arg.lineno,
                'col': arg.col_offset
            }
            
            if i < len(args.kw_defaults) and args.kw_defaults[i]:
                arg_info['default'] = self._ast_to_string(args.kw_defaults[i])
            
            arguments.append(arg_info)
        
        # **kwargs
        if args.kwarg:
            arguments.append({
                'name': args.kwarg.arg,
                'type': 'kwarg',
                'annotation': self._get_annotation(args.kwarg.annotation) if args.kwarg.annotation else None,
                'line': args.kwarg.lineno,
                'col': args.kwarg.col_offset
            })
        
        return arguments
    
    def _get_annotation(self, annotation: ast.AST) -> str:
        """Get string representation of type annotation."""
        return self._ast_to_string(annotation)
    
    def _ast_to_string(self, node: ast.AST) -> str:
        """Convert AST node to string representation."""
        try:
            return ast.unparse(node)
        except:
            # Fallback for older Python versions
            if isinstance(node, ast.Name):
                return node.id
            elif isinstance(node, ast.Constant):
                return repr(node.value)
            elif isinstance(node, ast.Attribute):
                return f"{self._ast_to_string(node.value)}.{node.attr}"
            else:
                return '<unknown>'
    
    def visit_FunctionDef(self, node: ast.FunctionDef):
        """Visit function definition."""
        old_function = self.current_function
        
        function_info = {
            'name': node.name,
            'line_start': node.lineno,
            'line_end': getattr(node, 'end_lineno', node.lineno),
            'col_start': node.col_offset,
            'col_end': getattr(node, 'end_col_offset', node.col_offset),
            'args': self._extract_args(node.args),
            'returns': self._get_annotation(node.returns) if node.returns else None,
            'decorators': [self._ast_to_string(dec) for dec in node.decorator_list],
            'docstring': self._get_docstring(node),
            'is_async': False,
            'is_method': self.current_class is not None,
            'class_name': self.current_class['name'] if self.current_class else None,
            'scope': self.current_scope['name'] if self.current_scope else 'global',
            'calls': [],  # Will be populated during call analysis
            'local_variables': []
        }
        
        self.functions.append(function_info)
        self.current_function = function_info
        
        # Enter function scope
        scope = self._enter_scope('function', node.name, node)
        
        # Process decorators
        for decorator in node.decorator_list:
            self.decorators.append({
                'target': node.name,
                'target_type': 'function',
                'decorator': self._ast_to_string(decorator),
                'line': decorator.lineno,
                'col': decorator.col_offset
            })
        
        # Visit function body
        self.generic_visit(node)
        
        # Exit function scope
        self._exit_scope()
        self.current_function = old_function
    
    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
        """Visit async function definition."""
        # Handle similar to regular function but mark as async
        old_function = self.current_function
        
        function_info = {
            'name': node.name,
            'line_start': node.lineno,
            'line_end': getattr(node, 'end_lineno', node.lineno),
            'col_start': node.col_offset,
            'col_end': getattr(node, 'end_col_offset', node.col_offset),
            'args': self._extract_args(node.args),
            'returns': self._get_annotation(node.returns) if node.returns else None,
            'decorators': [self._ast_to_string(dec) for dec in node.decorator_list],
            'docstring': self._get_docstring(node),
            'is_async': True,
            'is_method': self.current_class is not None,
            'class_name': self.current_class['name'] if self.current_class else None,
            'scope': self.current_scope['name'] if self.current_scope else 'global',
            'calls': [],
            'local_variables': []
        }
        
        self.functions.append(function_info)
        self.current_function = function_info
        
        scope = self._enter_scope('async_function', node.name, node)
        self.generic_visit(node)
        self._exit_scope()
        self.current_function = old_function
    
    def visit_ClassDef(self, node: ast.ClassDef):
        """Visit class definition."""
        old_class = self.current_class
        
        class_info = {
            'name': node.name,
            'line_start': node.lineno,
            'line_end': getattr(node, 'end_lineno', node.lineno),
            'col_start': node.col_offset,
            'col_end': getattr(node, 'end_col_offset', node.col_offset),
            'bases': [self._ast_to_string(base) for base in node.bases],
            'decorators': [self._ast_to_string(dec) for dec in node.decorator_list],
            'docstring': self._get_docstring(node),
            'methods': [],
            'attributes': [],
            'scope': self.current_scope['name'] if self.current_scope else 'global'
        }
        
        self.classes.append(class_info)
        self.current_class = class_info
        
        # Enter class scope
        scope = self._enter_scope('class', node.name, node)
        
        # Process decorators
        for decorator in node.decorator_list:
            self.decorators.append({
                'target': node.name,
                'target_type': 'class',
                'decorator': self._ast_to_string(decorator),
                'line': decorator.lineno,
                'col': decorator.col_offset
            })
        
        # Visit class body
        self.generic_visit(node)
        
        # Extract methods from functions found in this class
        class_methods = [f for f in self.functions if f.get('class_name') == node.name]
        class_info['methods'] = class_methods
        
        self._exit_scope()
        self.current_class = old_class
    
    def visit_Import(self, node: ast.Import):
        """Visit import statement."""
        for alias in node.names:
            import_info = {
                'type': 'import',
                'module': alias.name,
                'name': alias.asname if alias.asname else alias.name,
                'alias': alias.asname,
                'line': node.lineno,
                'col': node.col_offset
            }
            self.imports.append(import_info)
    
    def visit_ImportFrom(self, node: ast.ImportFrom):
        """Visit from-import statement."""
        module = node.module if node.module else ''
        level = node.level
        
        for alias in node.names:
            import_info = {
                'type': 'from_import',
                'module': module,
                'name': alias.name,
                'alias': alias.asname,
                'level': level,
                'line': node.lineno,
                'col': node.col_offset
            }
            self.imports.append(import_info)
    
    def visit_Call(self, node: ast.Call):
        """Visit function call."""
        call_info = {
            'line': node.lineno,
            'col': node.col_offset,
            'args': len(node.args),
            'kwargs': len(node.keywords),
            'function': self.current_function['name'] if self.current_function else None,
            'class': self.current_class['name'] if self.current_class else None,
            'scope': self.current_scope['name'] if self.current_scope else 'global'
        }
        
        # Determine call type and target
        if isinstance(node.func, ast.Name):
            # Simple function call
            call_info['type'] = 'function_call'
            call_info['target'] = node.func.id
            self.function_calls.append(call_info)
        elif isinstance(node.func, ast.Attribute):
            # Method call or attribute access
            call_info['type'] = 'method_call'
            call_info['target'] = node.func.attr
            call_info['object'] = self._ast_to_string(node.func.value)
            self.method_calls.append(call_info)
        else:
            # Complex call (e.g., function returned by another function)
            call_info['type'] = 'complex_call'
            call_info['target'] = self._ast_to_string(node.func)
            self.function_calls.append(call_info)
        
        # Add to current function's calls if we're in a function
        if self.current_function:
            self.current_function['calls'].append(call_info)
        
        self.generic_visit(node)
    
    def visit_Attribute(self, node: ast.Attribute):
        """Visit attribute access."""
        if not isinstance(node.ctx, ast.Store):  # Only record attribute reads
            attr_info = {
                'object': self._ast_to_string(node.value),
                'attribute': node.attr,
                'line': node.lineno,
                'col': node.col_offset,
                'function': self.current_function['name'] if self.current_function else None,
                'class': self.current_class['name'] if self.current_class else None,
                'scope': self.current_scope['name'] if self.current_scope else 'global'
            }
            self.attribute_access.append(attr_info)
        
        self.generic_visit(node)
    
    def visit_Assign(self, node: ast.Assign):
        """Visit assignment statement."""
        # Extract variable assignments
        for target in node.targets:
            if isinstance(target, ast.Name):
                var_info = {
                    'name': target.id,
                    'value': self._ast_to_string(node.value),
                    'line': node.lineno,
                    'col': node.col_offset,
                    'function': self.current_function['name'] if self.current_function else None,
                    'class': self.current_class['name'] if self.current_class else None,
                    'scope': self.current_scope['name'] if self.current_scope else 'global'
                }
                
                if self.current_function:
                    self.current_function['local_variables'].append(var_info)
                elif self.current_class and not self.current_function:
                    self.current_class['attributes'].append(var_info)
                else:
                    self.global_variables.append(var_info)
        
        self.generic_visit(node)
    
    def visit_Constant(self, node: ast.Constant):
        """Visit constant value."""
        if isinstance(node.value, (str, int, float, bool)) and node.value is not None:
            const_info = {
                'value': node.value,
                'type': type(node.value).__name__,
                'line': node.lineno,
                'col': node.col_offset,
                'function': self.current_function['name'] if self.current_function else None,
                'class': self.current_class['name'] if self.current_class else None,
                'scope': self.current_scope['name'] if self.current_scope else 'global'
            }
            self.constants.append(const_info)
        
        self.generic_visit(node)